#pragma once

#include <cassert>
#include <optional>
#include <limits>
#include <vector>
#include <initializer_list>
#include <utility>
#include <type_traits>

#include "txdnn_internal.h"

#include "operators/names.h"
#include "operators/invoker_params_base.h"

#include "operators/operator_struct.h"
#include "operators/operator_base.h"


namespace txdnn {
namespace operators {

template <class F, class... Ts>
void each_args(F f, Ts&&... xs) {
  (void)std::initializer_list<int>{(f(std::forward<Ts>(xs)), 0)...};
}


template <class Operator, class Context, class OperatorDesc>
Operators FindOperator(Operator op,
                       const Context& ctx,
                       const OperatorDesc& desc,
                       const AnyInvokeParams& invoke_ctx) {
  static_assert(sizeof(op) == sizeof(OperatorBase), "Operator must be stateless");
  static_assert(std::is_base_of<OperatorBase, Operator>{}, "Operator must inherit from OperatorBase");
  return op.getOperator(ctx, desc);
}

template <class... Operations>
struct OperatorContainer {

  template <class... OperatorsRight>
  auto operator+(OperatorContainer<OperatorsRight...>) const {
    return OperatorContainer<Operations..., OperatorsRight...>{};
  }

  // TODO
  template <class OperatorDesc, class Ops = txdnn::operators::Operators>
  std::vector<Ops> searchForOperators(const ExecutionContext& ctx,
                                      const OperatorDesc& desc,
                                      std::size_t limit = std::numeric_limits<std::size_t>::max(),
                                      const AnyInvokeParams& invoke_params = {}) const {
    std::vector<Ops> ops;
    std::size_t count = 0;
    each_args(
      [&](auto op) {
        if (count >= limit) { return; }
        {
          auto o = FindOperator(op, ctx, desc, invoke_params);
          if (o.succeeded()) {
            ++count;
            ops.emplace_back(std::move(o));
          } else {
            // TODO TXLOG_DEBUG();
          }
        }
      },
      Operations{}...
    );
    return ops;
  }

  // template <class OperatorDesc>
  // void ExecutePrimitive(const ExecutionContext& ctx,
  //                       const OperatorDesc& desc,
  //                       const AlgorithmName& name,
  //                       const AnyInvokeParams& invoke_params) const {
  //   const auto cfg = desc.makeNetworkConfig();
  //   if (const auto existing_invoker = ctx.getStream().getInvoker(cfg, std::nullopt, name)) {
  //     (*existing_invoker)(ctx.getStream(), invoke_params);
  //     return;
  //   }

  //   const auto ops = searchForOperators(ctx, desc, 1, invoke_params);
  //   if (ops.empty()) {
  //     TXDNN_THROW("No operator found");
  //   }

  //   const auto& op = ops.front();
  //   if (!op.invoker_factory_) {
  //     TXDNN_THROW("Operator has no invoker");
  //   }
  //   const auto invoker = ctx.getStream().prepareInvoker(*op.invoker_factory_, op.construct_params_); 
  //   ctx.getStream().registerInvoker(invoker, cfg, op.op_id_, name);
  //   invoker(ctx.getStream(), invoke_params);
  // }


  template <class OperatorDesc>
  void ExecutePrimitive(Handle& handle,
                        const OperatorDesc& desc,
                        const AlgorithmName& name,
                        const AnyInvokeParams& invoke_params) const {
    const auto cfg = desc.makeNetworkConfig();

    if (const auto existing_invoker = handle.getInvoker(cfg, std::nullopt, name)) {
      (*existing_invoker)(handle, invoke_params);
      return;
    }

    auto ctx = ExecutionContext{&handle};

    const auto ops = searchForOperators(ctx, desc, 1, invoke_params);
    if (ops.empty()) {
      TXDNN_THROW("No operator found");
    }

    const auto& op = ops.front();
    if (!op.invoker_factory_) {
      TXDNN_THROW("Operator has no invoker");
    }
    auto invoker = handle.prepareInvoker(*op.invoker_factory_, op.construct_params_); 
    handle.registerInvoker(invoker, cfg, op.op_id_, name);
    invoker(handle, invoke_params);
  }


  // template <class OperatorDesc>
  // void ExecutePrimitive(const Handle& handle,
  //                       const OperatorDesc& desc,
  //                       const AlgorithmName& name,
  //                       const AnyInvokeParams& invoke_params) const {
  //   return ExecutePrimitive(&handle, desc, name, invoke_params);
  // }

};

} // namespace operators
} // namespace txdnn